#!/usr/bin/env python3
"""
数据库操作模块 - PostgreSQL
本模块被 consumer.py（本地测试）和 lambda_handler.py（AWS部署）共享
"""

import os
import logging
from typing import Dict, Any, List
import psycopg2
from psycopg2.extras import Json

logger = logging.getLogger(__name__)


class PostgresWriter:
    """
    PostgreSQL 写入器

    功能：
    - 连接管理
    - 表初始化
    - 批量写入
    - 单条写入
    """

    def __init__(self, host=None, port=None, database=None, user=None, password=None):
        """
        初始化数据库连接配置

        Args:
            host: 数据库主机（默认从环境变量读取）
            port: 数据库端口
            database: 数据库名
            user: 用户名
            password: 密码
        """
        self.host = host or os.getenv('POSTGRES_HOST', 'postgres')
        self.port = port or int(os.getenv('POSTGRES_PORT', '5432'))
        self.database = database or os.getenv('POSTGRES_DB', 'behavior_analysis')
        self.user = user or os.getenv('POSTGRES_USER', 'postgres')
        self.password = password or os.getenv('POSTGRES_PASSWORD', 'postgres')

        self.conn = None
        self.cursor = None

    def connect(self):
        """连接数据库"""
        try:
            self.conn = psycopg2.connect(
                host=self.host,
                port=self.port,
                database=self.database,
                user=self.user,
                password=self.password
            )
            self.cursor = self.conn.cursor()
            logger.info(f"✅ 成功连接到PostgreSQL: {self.host}:{self.port}/{self.database}")

            # 创建表（如果不存在）
            self._create_table()
        except Exception as e:
            logger.error(f"❌ 连接PostgreSQL失败: {e}")
            raise

    def _create_table(self):
        """创建表（从 schema.sql 读取）"""
        try:
            # 读取 schema.sql 文件
            schema_path = os.path.join(os.path.dirname(__file__), 'schema.sql')
            with open(schema_path, 'r', encoding='utf-8') as f:
                schema_sql = f.read()

            # 执行 SQL（分号分隔的多个语句）
            # 移除注释和空行
            sql_statements = []
            for statement in schema_sql.split(';'):
                # 移除注释
                lines = []
                for line in statement.split('\n'):
                    # 移除 -- 注释
                    if '--' in line:
                        line = line[:line.index('--')]
                    lines.append(line)
                statement = '\n'.join(lines).strip()

                if statement:  # 非空语句
                    sql_statements.append(statement)

            # 执行所有 SQL 语句
            for sql in sql_statements:
                self.cursor.execute(sql)
            self.conn.commit()
            logger.info("✅ 数据库表已准备就绪")

        except FileNotFoundError:
            logger.error("❌ schema.sql 文件未找到")
            raise
        except Exception as e:
            logger.warning(f"创建表时出现警告（可能已存在）: {e}")
            self.conn.rollback()

    def write_behavior(self, event_data: Dict[str, Any], analysis_result: Dict[str, Any]):
        """
        写入单条行为分析结果

        Args:
            event_data: 原始事件数据
            analysis_result: 分析结果
        """
        try:
            insert_sql = """
            INSERT INTO user_behaviors
            (user_id, session_id, behavior_type, behavior_description, url,
             confidence, analysis_method, raw_event)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """

            self.cursor.execute(insert_sql, (
                event_data.get('userId', 'unknown'),
                event_data.get('sessionId', ''),
                analysis_result['behavior_type'],
                analysis_result['behavior_description'],
                event_data.get('url', ''),
                analysis_result.get('confidence', 0.0),
                analysis_result.get('analysis_method', 'unknown'),
                Json(event_data)
            ))

            self.conn.commit()
            logger.debug(f"✅ 写入行为记录: {analysis_result['behavior_description'][:50]}...")
        except Exception as e:
            logger.error(f"❌ 写入PostgreSQL失败: {e}")
            self.conn.rollback()
            raise

    def write_behaviors_batch(self, records: List[tuple]):
        """
        批量写入行为分析结果（性能优化）

        Args:
            records: [(event_data, analysis_result), ...]
        """
        if not records:
            return

        try:
            insert_sql = """
            INSERT INTO user_behaviors
            (user_id, session_id, behavior_type, behavior_description, url,
             confidence, analysis_method, raw_event)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """

            # 准备批量数据
            batch_data = []
            for event_data, analysis_result in records:
                batch_data.append((
                    event_data.get('userId', 'unknown'),
                    event_data.get('sessionId', ''),
                    analysis_result['behavior_type'],
                    analysis_result['behavior_description'],
                    event_data.get('url', ''),
                    analysis_result.get('confidence', 0.0),
                    analysis_result.get('analysis_method', 'unknown'),
                    Json(event_data)
                ))

            # 批量执行
            self.cursor.executemany(insert_sql, batch_data)
            self.conn.commit()
            logger.info(f"✅ 批量写入 {len(records)} 条行为记录")
        except Exception as e:
            logger.error(f"❌ 批量写入PostgreSQL失败: {e}")
            self.conn.rollback()
            raise

    def close(self):
        """关闭连接"""
        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()
        logger.debug("数据库连接已关闭")


def write_to_postgres(event_data: Dict[str, Any], analysis_result: Dict[str, Any]):
    """
    便捷函数：写入单条记录（Lambda 使用）

    这个函数每次都创建新连接，适合 Lambda 短生命周期场景。
    对于长时间运行的服务（consumer.py），推荐使用 PostgresWriter 类。

    Args:
        event_data: 原始事件数据
        analysis_result: 分析结果
    """
    writer = PostgresWriter()
    try:
        writer.connect()
        writer.write_behavior(event_data, analysis_result)
    finally:
        writer.close()


def write_to_postgres_batch(records: List[tuple]):
    """
    便捷函数：批量写入记录（Lambda 批处理使用）

    Args:
        records: [(event_data, analysis_result), ...]
    """
    writer = PostgresWriter()
    try:
        writer.connect()
        writer.write_behaviors_batch(records)
    finally:
        writer.close()
